//
// Copyright by mentu.zhou@outlook.com, all right reserved.
//


#ifdef __cplusplus
extern "C" {
#endif // __cpusplus

#include <ndis.h>

#ifdef __cplusplus
}
#endif // __cpusplus

#include "ndis_send_hook.h"
#include "dasm.h"
#include "code_protection.h"
#include "packet_processing.h"

#define NDIS_HOOK_TAG 'SIDN'

typedef struct _CURRUPTTED_FUNCTIONS
{
    LIST_ENTRY list_entry;
    PNDIS_OPEN_BLOCK ndis_block;
    SEND_PACKETS_HANDLER currupted_function;
    SEND_PACKETS_HANDLER repaired_function;
}CURRUPTTED_FUNCTIONS, *PCURRUPTTED_FUNCTIONS;

static BOOLEAN m_is_hooked = FALSE;
static ULONG m_hook_return_address[4] = {0};
static PUCHAR m_hook_entrys = NULL;

static LIST_ENTRY m_currupted_function_list = 
{
    &m_currupted_function_list, 
    &m_currupted_function_list
};

static void _hook_agent();

static VOID __stdcall _hook_calling_w_send_packets_handle(
    PNDIS_OPEN_BLOCK OpenBlock,
    NDIS_HANDLE MiniportAdapterContext,
    PPNDIS_PACKET PacketArray, 
    UINT NumberOfPackets
    );
static BOOLEAN _make_detour_over_hooked_entry(
    PNDIS_OPEN_BLOCK OpenBlock
    );

NTSTATUS ndis_hook_send_x(
    IN PNDIS_OPEN_BLOCK NdisOpenBlock
    )
{
#define ExitFunction goto __exit_function
    NTSTATUS status = STATUS_SUCCESS;
    PUCHAR ndis_m_send_x = NULL;
    PUCHAR current_op_in_send_x = NULL;
    ULONG current_op_length = 0;
    ULONG parsed_code_length = 0;
    ULONG old_status = 0;
    ULONG hook_pointer_count = 0;
    PULONG hook_entry_long = NULL;
    PLARGE_INTEGER hook_entry_large_int = NULL;
    UCHAR register_status = 0;
    PUCHAR protect_data = NULL;
    ULONG protect_size = 0;

    KdPrint(("work in : 0x%08X\r\n", NdisOpenBlock->NdisCommonOpenBlock.SendHandler));

    if (!m_is_hooked)
    {
        // Search for hook point.
        ndis_m_send_x = (PUCHAR)NdisOpenBlock->NdisCommonOpenBlock.SendHandler;
        current_op_in_send_x = ndis_m_send_x;

        // find first long distance JNE
        while (parsed_code_length < 0x100)
        {
            current_op_length = GetOpCodeSize(current_op_in_send_x);
            if (   current_op_length == 6
                && current_op_in_send_x[0] == 0x0f
                && current_op_in_send_x[1] == 0x85
                && *(PULONG)(current_op_in_send_x + 2) > 0x00001000)
            {
                current_op_in_send_x += *(PULONG)(current_op_in_send_x + 2);
                current_op_in_send_x += 6;
                break;
            }
            else
            {
                current_op_in_send_x += current_op_length;
                parsed_code_length += current_op_length;
            }
        }

        if (parsed_code_length >= 0x100)
        {
            KdPrint(("Fail to find long jump\r\n"));
            status = STATUS_UNSUCCESSFUL;
            ExitFunction;
        }

        parsed_code_length = 0;
        while (parsed_code_length < 0x100)
        {
            current_op_length = GetOpCodeSize(current_op_in_send_x);
            if (current_op_length == 6
                && current_op_in_send_x[0] == 0xff
                && *(PULONG)(current_op_in_send_x + 2) == 0x00000090)
            {
                break;
            }
            else
            {
                current_op_in_send_x += current_op_length;
                parsed_code_length += current_op_length;
            }
        }

        // Not found.
        if (parsed_code_length >= 0x300)
        {
            KdPrint(("Fail to find call entry\r\n"));
            status = STATUS_UNSUCCESSFUL;
            ExitFunction;
        }

        //
        // The first hook pointer is found.
        //
        parsed_code_length = 0;
        m_hook_return_address[0] = (ULONG)current_op_in_send_x;
        current_op_in_send_x += 6;
        hook_pointer_count = 1;
        while (parsed_code_length < 0x17000)
        {
            current_op_length = GetOpCodeSize(current_op_in_send_x);
            if (current_op_length == 6
                && current_op_in_send_x[0] == 0xff
                && *(PULONG)(current_op_in_send_x + 2) == 0x00000090)
            {
                m_hook_return_address[hook_pointer_count] = 
                    (ULONG)current_op_in_send_x;
                current_op_in_send_x += current_op_length;
                ++ hook_pointer_count;
                if (hook_pointer_count > 3)
                {
                    break;
                }
            }
            else
            {
                current_op_in_send_x += current_op_length;
            }
            parsed_code_length += current_op_length;
        }

        if (hook_pointer_count <= 3)
        {
            KdPrint(("Not enough hook pointer!\r\n"));
            status = STATUS_NOT_IMPLEMENTED;
            ExitFunction;
        }

        //
        // Found, hook it.
        //

        m_hook_entrys = (PUCHAR)ExAllocatePoolWithTag(
            NonPagedPool,
            8 * 4,
            NDIS_HOOK_TAG);
        hook_entry_long = (PULONG)m_hook_entrys;
        hook_entry_large_int = (PLARGE_INTEGER)(m_hook_entrys + 3);
        if (hook_entry_long == NULL)
        {
            KdPrint(("Fail to allocate entrys.\r\n"));
            status = STATUS_INSUFFICIENT_RESOURCES;
            ExitFunction;
        }

        hook_entry_long[0] = 0xc368c08b;
        hook_entry_long[1] = 0xc368c08b;
        hook_entry_long[2] = 0xc368c18b;
        hook_entry_long[3] = 0xc368c18b;
        hook_entry_long[4] = 0xc368c28b;
        hook_entry_long[5] = 0xc368c28b;
        hook_entry_long[6] = 0xc368c38b;
        hook_entry_long[7] = 0xc368c38b;

        protect_data = (PUCHAR)_hook_agent;
        hook_entry_large_int[0].LowPart = 
        hook_entry_large_int[1].LowPart = 
        hook_entry_large_int[2].LowPart = 
        hook_entry_large_int[3].LowPart = (ULONG)protect_data;

        protect_size = 0;
        while (protect_data[protect_size] != 0xcc)
        {
            protect_size += GetOpCodeSize(protect_data + protect_size);
        }
        ++ protect_size;
        add_protection(protect_data, protect_size);
        add_protection((PUCHAR)hook_entry_long, 32);

        protect_data = (PUCHAR)_hook_calling_w_send_packets_handle;
        protect_size = 0;
        while (protect_data[protect_size] != 0xcc)
        {
            protect_size += GetOpCodeSize(protect_data + protect_size);
        }
        ++ protect_size;

        // Disable protect.
        __asm
        {
            cli;
            mov eax, cr0;
            mov old_status, eax;
            and eax, 0xfffeffff;
            mov cr0, eax;
        }

        current_op_in_send_x = (PUCHAR)m_hook_return_address[0];
        register_status = current_op_in_send_x[1];
        *(PUSHORT)current_op_in_send_x = 0xe890;
        *(PULONG)(current_op_in_send_x + 2) = 
            (ULONG)hook_entry_long 
            + 8 * (register_status - 0x90)
            - ((ULONG)current_op_in_send_x + 6);
        KdPrint(("Hook at %08X\r\n", current_op_in_send_x));

        current_op_in_send_x = (PUCHAR)m_hook_return_address[1];
        register_status = current_op_in_send_x[1];
        *(PUSHORT)current_op_in_send_x = 0xe890;
        *(PULONG)(current_op_in_send_x + 2) = 
            (ULONG)hook_entry_long 
            + 8 * (register_status - 0x90)
            - ((ULONG)current_op_in_send_x + 6);
        KdPrint(("Hook at %08X\r\n", current_op_in_send_x));

        current_op_in_send_x = (PUCHAR)m_hook_return_address[2];
        register_status = current_op_in_send_x[1];
        *(PUSHORT)current_op_in_send_x = 0xe890;
        *(PULONG)(current_op_in_send_x + 2) = 
            (ULONG)hook_entry_long 
            + 8 * (register_status - 0x90)
            - ((ULONG)current_op_in_send_x + 6);
        KdPrint(("Hook at %08X\r\n", current_op_in_send_x));

        current_op_in_send_x = (PUCHAR)m_hook_return_address[3];
        register_status = current_op_in_send_x[1];
        *(PUSHORT)current_op_in_send_x = 0xe890;
        *(PULONG)(current_op_in_send_x + 2) = 
            (ULONG)hook_entry_long 
            + 8 * (register_status - 0x90)
            - ((ULONG)current_op_in_send_x + 6);
        KdPrint(("Hook at %08X\r\n", current_op_in_send_x));

        // Enable protect.
        __asm
        {
            mov eax, old_status;
            mov cr0, eax;
            sti;
        }

        for (int i = 0; i < 4; ++ i)
        {
            add_protection((PUCHAR)m_hook_return_address[i] + 1, 5);
        }

        m_is_hooked = TRUE;
    }

    //
    // make a detour function to avoid the hooked entry.
    //
    _make_detour_over_hooked_entry(NdisOpenBlock);

    return STATUS_SUCCESS;

__exit_function:
    return status;
}

NTSTATUS ndis_unhook_send_x()
{
    PUCHAR hooked_address = 0;
    ULONG old_status = 0;
    PLIST_ENTRY list_entry = NULL;
    PVOID buffer = NULL;

    if (!m_is_hooked)
    {
        return STATUS_SUCCESS;
    }

    // Disable protect.
    __asm
    {
        cli;
        mov eax, cr0;
        mov old_status, eax;
        and eax, 0xfffeffff;
        mov cr0, eax;
    }

    hooked_address = (PUCHAR)m_hook_return_address[0];
    hooked_address[0] = 0xFF;
    hooked_address[1] = UCHAR((hooked_address + 6 + *(PULONG)(hooked_address + 2) - m_hook_entrys) / 8) + 0x90;
    *(PULONG)(hooked_address + 2) = 0x00000090;

    hooked_address = (PUCHAR)m_hook_return_address[1];
    hooked_address[0] = 0xFF;
    hooked_address[1] = UCHAR((hooked_address + 6 + *(PULONG)(hooked_address + 2) - m_hook_entrys) / 8) + 0x90;
    *(PULONG)(hooked_address + 2) = 0x00000090;

    hooked_address = (PUCHAR)m_hook_return_address[2];
    hooked_address[0] = 0xFF;
    hooked_address[1] = UCHAR((hooked_address + 6 + *(PULONG)(hooked_address + 2) - m_hook_entrys) / 8) + 0x90;
    *(PULONG)(hooked_address + 2) = 0x00000090;

    hooked_address = (PUCHAR)m_hook_return_address[3];
    hooked_address[0] = 0xFF;
    hooked_address[1] = UCHAR((hooked_address + 6 + *(PULONG)(hooked_address + 2) - m_hook_entrys) / 8) + 0x90;
    *(PULONG)(hooked_address + 2) = 0x00000090;

    // Enable protect.
    __asm
    {
        mov eax, old_status;
        mov cr0, eax;
        sti;
    }

    ExFreePoolWithTag(m_hook_entrys, NDIS_HOOK_TAG);

    //
    // Release all allocated buffer.
    //
    for (list_entry = m_currupted_function_list.Flink; 
         list_entry != &m_currupted_function_list;
         )
    {
        buffer = (PVOID)list_entry;
        list_entry = list_entry->Flink;
        ExFreePoolWithTag(buffer, NDIS_HOOK_TAG);
    }

    return STATUS_SUCCESS;
}

__declspec(naked) void _hook_agent()
{
#define miniport_handle (04h)
#define packet_list (08h)
#define packet_count (0ch)

    __asm
    {
        mov ecx, dword ptr [esp+packet_count];
        mov edx, dword ptr [esp+packet_list];
        mov ebx, dword ptr [esp+miniport_handle];
        push ecx;
        push edx;
        push ebx;
        push eax;
        call _hook_calling_w_send_packets_handle;
        ret 0ch;
    }
}

void __stdcall _hook_calling_w_send_packets_handle(
    PNDIS_OPEN_BLOCK OpenBlock,
    NDIS_HANDLE MiniportAdapterContext, 
    PPNDIS_PACKET PacketArray, 
    UINT NumberOfPackets 
    )
{
    PLIST_ENTRY list_entry = NULL;
    PCURRUPTTED_FUNCTIONS curruptted_function = NULL;
    BOOLEAN is_found = FALSE;
    UINT index = 0;
    SEND_PACKETS_HANDLER send_packets_handler = NULL;

    send_packets_handler = (SEND_PACKETS_HANDLER)((PULONG)OpenBlock)[0x24];

    for (list_entry = m_currupted_function_list.Flink; 
        list_entry != &m_currupted_function_list; 
        list_entry = list_entry->Flink)
    {
        curruptted_function = PCURRUPTTED_FUNCTIONS(list_entry);
        if (   OpenBlock == curruptted_function->ndis_block
            || send_packets_handler == curruptted_function->currupted_function)
        {
            is_found = TRUE;
        }
    }

    if (!is_found)
    {
        return send_packets_handler(
            MiniportAdapterContext,
            PacketArray,
            NumberOfPackets);
    }

    for (index = 0; index < NumberOfPackets; ++ index)
    {
        packet_processing(PacketArray[index]);
    }

    return curruptted_function->repaired_function(
        MiniportAdapterContext,
        PacketArray,
        NumberOfPackets);
}

BOOLEAN _make_detour_over_hooked_entry(
    PNDIS_OPEN_BLOCK OpenBlock
    )
{
    //
    // In different condition, need to repair it in many ways.
    // Here, it's an example, just copy original code to make an detour.
    //
    
    PLIST_ENTRY list_entry = NULL;
    PCURRUPTTED_FUNCTIONS curruptted_function = NULL;
    ULONG code_size = 0;
    PUCHAR detour_function = NULL;
    PUCHAR enemy_hook_agent = NULL;
    PUCHAR enemy_restored_head = NULL;
    PUCHAR enemy_jmp_original = NULL;
    UCHAR code_buffer[32] = {0};

    PUCHAR hooked_function = (PUCHAR)(((PULONG)OpenBlock)[0x24]);

    for (list_entry = m_currupted_function_list.Flink; 
         list_entry != &m_currupted_function_list; 
         list_entry = list_entry->Flink)
    {
        curruptted_function = PCURRUPTTED_FUNCTIONS(list_entry);
        if (hooked_function == (PUCHAR)curruptted_function->currupted_function
            && OpenBlock == curruptted_function->ndis_block)
        {
            return TRUE;
        }
    }

    // Keep function header in buffer in case of sudden hook.
    __asm
    {
        cli;
    }
    RtlCopyMemory(code_buffer, hooked_function, 32);
    __asm
    {
        sti;
    }

    // 
    // If begin with 'JMP XXXXXXXX'
    // Assume it's hooked, find the original code header.
    // Make a detour.
    //

    if (code_buffer[0] == 0xe9)
    {
        // Calculate the hook agent address.
        enemy_hook_agent = hooked_function + 5 + *(PULONG)(code_buffer + 1);
        KdPrint(("Hooked fucntion : 0x%08X=>0x%08X\r\n", hooked_function, enemy_hook_agent));

        // Search for jumping to place where keep the original head.
        code_size = 0;
        while (code_size < 200)
        {
            if (enemy_hook_agent[code_size] == 0xe9)
            {
                break;
            }
            code_size += GetOpCodeSize(enemy_hook_agent + code_size);
        }
        if (code_size >= 200)
        {
            KdPrint(("Failed to find jumping to place where keep the original header.\r\n"));
            return FALSE;
        }
        enemy_restored_head = enemy_hook_agent 
                              + code_size 
                              + 5
                              + *(PULONG)(enemy_hook_agent + code_size + 1);
        KdPrint(("Origina head stored at : 0x%08X\r\n", enemy_restored_head));

        // Search for jumping to original body. Get size of head.
        code_size = 0;
        while (code_size < 32)
        {
            if (enemy_restored_head[code_size] == 0xe9)
            {
                break;
            }
            code_size += GetOpCodeSize(enemy_restored_head + code_size);
        }
        if (code_size >= 32)
        {
            KdPrint(("Origina head stored seem to be bad.\r\n"));
            return FALSE;
        }
        enemy_jmp_original = enemy_restored_head
                             + code_size
                             + 5
                             + *(PULONG)(enemy_restored_head + code_size + 1);

        // Vilify it.
        if (hooked_function + code_size != enemy_jmp_original)
        {
            KdPrint(("Origina head stored fail to vilify.\r\n"));
            return FALSE;
        }
    }
    else
    {
        code_size = 0;
        while (code_size < 5)
        {
            code_size += GetOpCodeSize(code_buffer + code_size);
        }
        enemy_restored_head = code_buffer;
        enemy_jmp_original = hooked_function + code_size;
    }

    curruptted_function = (PCURRUPTTED_FUNCTIONS)ExAllocatePoolWithTag(
        NonPagedPool, 
        sizeof(CURRUPTTED_FUNCTIONS) + code_size + 6,
        NDIS_HOOK_TAG);
    if (curruptted_function == NULL)
    {
        return FALSE;
    }

    detour_function = (PUCHAR)(curruptted_function + 1);

    curruptted_function->ndis_block = OpenBlock;
    curruptted_function->currupted_function 
        = (SEND_PACKETS_HANDLER)hooked_function;
    curruptted_function->repaired_function 
        = (SEND_PACKETS_HANDLER)detour_function;

    RtlCopyMemory(detour_function, enemy_restored_head, code_size);
    detour_function[code_size] = 0x68;
    *(PULONG)(detour_function + code_size + 1) = (ULONG)enemy_jmp_original;
    detour_function[code_size + 5] = 0xc3;
    InsertTailList(
        &m_currupted_function_list, 
        &curruptted_function->list_entry);
    add_protection(
        (PUCHAR)&curruptted_function->ndis_block,
        sizeof(CURRUPTTED_FUNCTIONS) + code_size - 2);

    return TRUE;
}
